import gym
from gym import spaces
import numpy as np
import torch
import torch.nn as nn
from torch.distributions import MultivariateNormal
import matplotlib.pyplot as plt
from collections import deque
import random

# 设置随机种子以确保可重复性
torch.manual_seed(0)
np.random.seed(0)


# 自定义无人机环境
class UAVEnv(gym.Env):
    def __init__(self):
        super(UAVEnv, self).__init__()

        # 环境参数
        self.area_size = 1000.0  # 区域大小 1000x1000 m
        self.num_users = 10  # 用户数量
        self.max_steps = 200  # 每个episode最大步数
        self.uav_speed = 20.0  # 无人机速度 m/s
        self.collection_range = 50.0  # 数据收集范围 m

        # 动作空间: [v_x, v_y] 速度向量 (-1 to 1)
        self.action_space = spaces.Box(
            low=-1.0, high=1.0, shape=(2,), dtype=np.float32)

        # 状态空间: 无人机位置 + 所有用户位置 + 用户数据状态
        # [drone_x, drone_y, user1_x, user1_y, user1_has_data, ...]
        self.observation_space = spaces.Box(
            low=0.0, high=self.area_size,
            shape=(2 + self.num_users * 3,), dtype=np.float32)

        # 重置环境
        self.reset()

    def reset(self):
        # 重置步数计数器
        self.steps = 0

        # 随机初始化无人机位置 (区域中心附近)
        self.drone_pos = np.array([self.area_size / 2, self.area_size / 2])

        # 随机初始化用户位置 (在整个区域)
        self.user_positions = np.random.rand(self.num_users, 2) * self.area_size

        # 初始化用户数据状态 (1表示有数据待收集)
        self.user_has_data = np.ones(self.num_users)

        # 返回初始状态
        return self._get_state()

    def _get_state(self):
        # 构建状态向量
        state = [self.drone_pos[0], self.drone_pos[1]]
        for i in range(self.num_users):
            state.extend([
                self.user_positions[i, 0],
                self.user_positions[i, 1],
                self.user_has_data[i]
            ])
        return np.array(state, dtype=np.float32)

    def step(self, action):
        # 增加步数
        self.steps += 1

        # 处理动作 (归一化到实际速度)
        norm_action = np.clip(action, -1.0, 1.0)
        movement = norm_action * self.uav_speed

        # 更新无人机位置
        new_pos = self.drone_pos + movement
        new_pos = np.clip(new_pos, 0.0, self.area_size)
        self.drone_pos = new_pos

        # 检查数据收集
        collected_data = 0
        for i in range(self.num_users):
            if self.user_has_data[i] > 0.5:  # 如果用户有数据
                distance = np.linalg.norm(self.drone_pos - self.user_positions[i])
                if distance < self.collection_range:
                    self.user_has_data[i] = 0  # 标记为已收集
                    collected_data += 1

        # 计算奖励
        reward = collected_data * 10  # 每收集一个数据点奖励10分

        # 添加小型惩罚鼓励高效移动
        reward -= 0.1 * np.linalg.norm(movement)

        # 检查是否结束
        done = (self.steps >= self.max_steps) or (np.sum(self.user_has_data) < 0.5)

        # 如果收集了所有数据，额外奖励
        if np.sum(self.user_has_data) < 0.5:
            reward += 100  # 完成所有收集的额外奖励

        # 返回状态、奖励、完成标志和信息
        return self._get_state(), reward, done, {}

    def render(self, mode='human'):
        # 简单可视化 (可选实现)
        pass


# 策略网络 (Actor-Critic结构)
class PolicyNetwork(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(PolicyNetwork, self).__init__()

        # 共享特征提取层
        self.shared_net = nn.Sequential(
            nn.Linear(state_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
        )

        # 策略头 (Actor)
        self.policy_head = nn.Sequential(
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, action_dim)
        )

        # 价值头 (Critic)
        self.value_head = nn.Sequential(
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 1)
        )

        # 动作标准差参数 (可学习)
        self.log_std = nn.Parameter(torch.zeros(1, action_dim))

    def forward(self, x):
        features = self.shared_net(x)
        action_mean = self.policy_head(features)
        value = self.value_head(features)
        return action_mean, value

    def act(self, state):
        state = torch.FloatTensor(state).unsqueeze(0)
        action_mean, value = self.forward(state)

        # 创建协方差矩阵
        action_std = torch.exp(self.log_std)
        cov_mat = torch.diag(action_std.squeeze(0)).unsqueeze(0)

        # 创建分布并采样
        dist = MultivariateNormal(action_mean, cov_mat)
        action = dist.sample()
        log_prob = dist.log_prob(action)

        return action.detach().numpy()[0], log_prob, value.detach().numpy()[0]


# PPO算法实现
class PPO:
    def __init__(self, state_dim, action_dim, lr=3e-4, gamma=0.99, eps_clip=0.2, K_epochs=10):
        self.gamma = gamma
        self.eps_clip = eps_clip
        self.K_epochs = K_epochs

        self.policy = PolicyNetwork(state_dim, action_dim)
        self.optimizer = torch.optim.Adam(self.policy.parameters(), lr=lr)

        self.MseLoss = nn.MSELoss()

    def update(self, memory):
        # 蒙特卡洛估计奖励
        rewards = []
        discounted_reward = 0
        for reward, is_terminal in zip(reversed(memory.rewards), reversed(memory.is_terminals)):
            if is_terminal:
                discounted_reward = 0
            discounted_reward = reward + (self.gamma * discounted_reward)
            rewards.insert(0, discounted_reward)

        # 归一化奖励
        rewards = torch.tensor(rewards, dtype=torch.float32)
        rewards = (rewards - rewards.mean()) / (rewards.std() + 1e-7)

        # 转换旧数据为张量
        old_states = torch.stack(memory.states).detach()
        old_actions = torch.stack(memory.actions).detach()
        old_logprobs = torch.stack(memory.logprobs).detach()

        # 多轮优化
        for _ in range(self.K_epochs):
            # 评估旧动作
            action_means, state_values = self.policy(old_states)
            action_std = torch.exp(self.policy.log_std)
            cov_mat = torch.diag(action_std.squeeze(0)).unsqueeze(0)
            dist = MultivariateNormal(action_means, cov_mat)

            # 找到新动作的log概率和熵
            logprobs = dist.log_prob(old_actions)
            dist_entropy = dist.entropy()

            # 计算比率 (pi_theta / pi_theta_old)
            ratios = torch.exp(logprobs - old_logprobs.detach())

            # 计算优势函数
            advantages = rewards - state_values.detach().squeeze()

            # 替代损失
            surr1 = ratios * advantages
            surr2 = torch.clamp(ratios, 1 - self.eps_clip, 1 + self.eps_clip) * advantages

            # 策略损失 + 价值损失 - 熵奖励
            loss = -torch.min(surr1, surr2) + 0.5 * self.MseLoss(state_values.squeeze(), rewards) - 0.01 * dist_entropy

            # 梯度下降
            self.optimizer.zero_grad()
            loss.mean().backward()
            self.optimizer.step()


# 存储轨迹数据的类
class Memory:
    def __init__(self):
        self.actions = []
        self.states = []
        self.logprobs = []
        self.rewards = []
        self.is_terminals = []

    def clear_memory(self):
        del self.actions[:]
        del self.states[:]
        del self.logprobs[:]
        del self.rewards[:]
        del self.is_terminals[:]


# 主训练函数
def main():
    # 创建环境
    env = UAVEnv()
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.shape[0]

    # 创建PPO代理
    ppo_agent = PPO(state_dim, action_dim)
    memory = Memory()

    # 训练参数
    max_episodes = 1000  # 减少这个值以快速测试
    update_timestep = 2000  # 更新策略的时间步
    timestep = 0
    i_episode = 0

    # 记录奖励
    reward_history = []
    avg_reward_history = []

    print("开始训练...")

    while i_episode < max_episodes:
        state = env.reset()
        episode_reward = 0

        for t in range(1, env.max_steps + 1):
            timestep += 1

            # 运行旧策略
            action, logprob, _ = ppo_agent.policy.act(state)

            # 执行动作
            next_state, reward, done, _ = env.step(action)

            # 存储转换
            memory.states.append(torch.FloatTensor(state))
            memory.actions.append(torch.FloatTensor(action))
            memory.logprobs.append(logprob)
            memory.rewards.append(reward)
            memory.is_terminals.append(done)

            # 更新状态
            state = next_state
            episode_reward += reward

            # 更新策略
            if timestep % update_timestep == 0:
                ppo_agent.update(memory)
                memory.clear_memory()
                timestep = 0

            if done:
                break

        # 记录奖励
        reward_history.append(episode_reward)

        # 计算平均奖励 (最近100个episode)
        if len(reward_history) > 100:
            avg_reward = np.mean(reward_history[-100:])
            avg_reward_history.append(avg_reward)

            # 打印进度
            if i_episode % 10 == 0:
                print(f"Episode: {i_episode}, Reward: {episode_reward:.2f}, Avg Reward: {avg_reward:.2f}")

                # 如果平均奖励足够高，可以提前停止
                if avg_reward > 180:  # 接近最大可能奖励
                    print("已收敛!")
                    break

        i_episode += 1

    # 绘制奖励曲线
    plt.figure(figsize=(10, 5))
    plt.plot(reward_history, alpha=0.3, label='每集奖励')
    plt.plot(avg_reward_history, label='平均奖励(100集)')
    plt.xlabel('集数')
    plt.ylabel('奖励')
    plt.title('训练进度')
    plt.legend()
    plt.grid(True)
    plt.savefig('training_rewards.png')
    plt.show()

    print("训练完成!")


if __name__ == '__main__':
    main()